import sys
from PyQt5.QtWidgets import QApplication, QWidget, QGridLayout, QLabel 
from PyQt5.QtGui import QPixmap
import sys
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *
import SimpleITK as sitk
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import cv2 

label_color = {
    1: [249,72,247],
    2: [209,198,103],
    3: [70,139,88],
    4: [199,104,19],
    5: [174,213,76],
    6: [83,83,50]
}

def ImageFusion(image, label):
    image_fusion = np.expand_dims(image, axis=-1) 
    image_fusion = np.repeat(image_fusion, 3, axis=-1)
    width = label.shape[1]
    height = label.shape[0]
    for i in range(width):
        for j in range(height):
            if label[j,i] == 0:
                continue
            
            image_fusion[j,i] = label_color[label[j,i]]
           
    return image_fusion
    

if __name__ == '__main__':
    image = sitk.ReadImage("/media/chyang/data/dataset/project/hipbone/nnUNet_raw/Dataset001_PelvicNAS/imagesTr/1.3.6.1.4.1.9328.50.4.0001_0000.nii.gz") #in_file是nii.gz文件的路径
    label = sitk.ReadImage("/media/chyang/data/dataset/project/hipbone/nnUNet_raw/Dataset001_PelvicNAS/labelsTr/1.3.6.1.4.1.9328.50.4.0001.nii.gz")
    image_arr = sitk.GetArrayFromImage(image)
    image_scaled = (image_arr - image_arr.min()) / (image_arr.max() - image_arr.min()) * 255
    image_scaled = image_scaled.astype(np.uint8)
    
    
    label_arr = sitk.GetArrayFromImage(label)
    label_scaled = label_arr
    label_scaled = label_scaled.astype(np.uint8)
    

    
    # z = 0
    # slice = img_scaled[:, :, 200]
    # plt.imshow(slice)
    # plt.show()
    # print("eee")
    
    # cv2.imshow( "image_slice",image_scaled[:, :, 200])
    # cv2.imshow( "label_slice",label_scaled[:, :, 200])
    image_fusion = ImageFusion(image_scaled[:, :, 200], label_scaled[:, :, 200])
    cv2.imshow( "image_fusion", image_fusion)
    cv2.waitKey(0)
    
    
    
    